import pandas as pd
from torch.utils.data import Dataset
from sklearn.preprocessing import LabelEncoder

from src.bert.data_preprocess import config


class TextDataset(Dataset):
    def __init__(self, dataframe):
        super().__init__()
        self.dataframe = dataframe

    def __getitem__(self, idx):
        text = self.dataframe.iloc[idx]['review_clean']
        # 限制文本长度为150个字符，超过部分截断
        if isinstance(text, str):
            text = text[:150]
        else:
            text = ""
        label = self.dataframe.iloc[idx]['label']
        cat = self.dataframe.iloc[idx]['cat']
        cat = cat - 1
        return text, label, cat

    def __len__(self):
        return len(self.dataframe)


def load_raw_data(file_path):
    df = pd.read_csv(file_path)
    selected_data = df[['review_clean', 'label', 'cat']]
    return selected_data


def cat_clean():
    # 读取数据
    df = pd.read_csv(config.test_path)

    # 标签编码
    encoder = LabelEncoder()
    df["cat_id"] = encoder.fit_transform(df["cat"])

    # 保存类别映射关系 encoder.classes_ 类别标签按字母顺序排列
    cat2id = dict(zip(encoder.classes_, encoder.transform(encoder.classes_)))  # [('手机', 0), ('水果', 1), ('衣服', 2)]

    # 用数字替换原来的 cat 列
    df["cat"] = df["cat_id"]
    df = df.drop(columns=["cat_id"])  # 删除临时列

    # 保存新的数据集
    df.to_csv(config.test_path, index=False, encoding="utf-8")

    print("类别映射:", cat2id)
    print("新数据集预览:\n", df.head())


if __name__ == '__main__':
    # data = load_raw_data(config.train_path)
    # print(data.head())
    # cat_clean()
    pass
